package utils;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Vector;

import java.util.regex.Pattern;

public class CrossValidator {
	HashMap<String, String> itemList = new HashMap<String, String>();
	HashMap<String, String> allItems = new HashMap<String, String>();

	/*
	 * Loads the list of files. The file list must have one file per line, with
	 * the file name being the first column and the label being the last column
	 */
	public void loadFiles(String itemListFile, boolean hasHeader) {
		FileInputStream fstream;
		try {
			fstream = new FileInputStream(itemListFile);
			BufferedReader in = new BufferedReader(new InputStreamReader(
					fstream));
			String line;
			if (hasHeader) {
				// Consume the header line first
				line = in.readLine();
			}
			do {
				line = in.readLine();
				if (line != null) {
					String regex = "\\s*,\\s*";
					Pattern p = Pattern.compile(regex);
					String results[] = p.split(line);
					itemList.put(results[0], results[results.length - 1]);
					allItems.put(results[0], line);
					System.out.println(results[0] + ":"
							+ results[results.length - 1]);
				}
			} while (line != null);
			System.out.println(itemList.size() + " files have been loaded.");
			fstream.close();
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} catch (NumberFormatException e) {
			e.printStackTrace();
		}
	}

	/*
	 * Implements the cross validation
	 * 
	 * @param numOfFolds The number of folds of cross validation by default the
	 * value is 10
	 * 
	 * @param location The location where the generated files are to be written
	 * to
	 */
	public void crossValidate(int numOfFolds, String location) {
		String trainingSets[] = new String[numOfFolds];
		String testingSets[] = new String[numOfFolds];
		// Prefill the instance strings
		for (int foldIndex = 0; foldIndex < numOfFolds; foldIndex++) {
			trainingSets[foldIndex] = "";
			testingSets[foldIndex] = "";
		}

		/*
		 * Randomise and stratify
		 */
		// Step 1: Determine the number of classes
		HashSet<String> uniqueLabels = new HashSet<String>(itemList.values());
		int numOfClasses = uniqueLabels.size();
		String[] labels = new String[numOfClasses];
		uniqueLabels.toArray(labels);

		// Step 2: Filter the items by class label and save them in separate
		// containers
		HashMap<String, Vector<String>> buckets = new HashMap<String, Vector<String>>();
		for (int i = 0; i < numOfClasses; i++) {
			// Prefill the buckets for each different classes
			buckets.put(labels[i], new Vector<String>());
		}

		for (String item : itemList.keySet()) {
			buckets.get(itemList.get(item)).add(item);
		}

		// Step 3: Randomise the instances in each container
		for (String label : buckets.keySet()) {
			Collections.shuffle(buckets.get(label));
		}

		/* Generate the training sets and testing sets */
		for (int classIndex = 0; classIndex < numOfClasses; classIndex++) {
			int instanceAnchors[] = new int[numOfFolds];// Saves the indices of
			// boundary items

			int numOfInstances = buckets.get(labels[classIndex]).size();
			System.out.println("Class " + labels[classIndex] + " has "
					+ numOfInstances + " instances.");

			for (int foldIndex = 0; foldIndex < numOfFolds; foldIndex++) {

				// Predetermine the boundaries of folds
				if (foldIndex == 0) {
					instanceAnchors[foldIndex] = numOfInstances / numOfFolds
							- 1;
				} else if (foldIndex == numOfFolds - 1) {
					instanceAnchors[foldIndex] = numOfInstances - 1;
				} else {
					instanceAnchors[foldIndex] = instanceAnchors[foldIndex - 1]
							+ numOfInstances / numOfFolds;
				}
				System.out.println("Fold " + (foldIndex + 1) + " ends at "
						+ instanceAnchors[foldIndex]);
			}

			// Retrieve the contents of data sets
			for (int instanceIndex = 0; instanceIndex < numOfInstances; instanceIndex++) {
				for (int foldIndex = 0; foldIndex < numOfFolds; foldIndex++) {
					if (foldIndex == 0) {
						if (instanceIndex >= 0
								&& instanceIndex < instanceAnchors[0]) {
							testingSets[foldIndex] = testingSets[foldIndex]
									.concat(allItems.get(buckets.get(
											labels[classIndex]).get(
											instanceIndex))
											+ "\n");
						} else {
							trainingSets[foldIndex] = trainingSets[foldIndex]
									.concat(allItems.get(buckets.get(
											labels[classIndex]).get(
											instanceIndex))
											+ "\n");
						}
					} else {
						if (instanceIndex >= instanceAnchors[foldIndex - 1]
								&& instanceIndex < instanceAnchors[foldIndex]) {
							testingSets[foldIndex] = testingSets[foldIndex]
									.concat(allItems.get(buckets.get(
											labels[classIndex]).get(
											instanceIndex))
											+ "\n");
						} else {
							trainingSets[foldIndex] = trainingSets[foldIndex]
									.concat(allItems.get(buckets.get(
											labels[classIndex]).get(
											instanceIndex))
											+ "\n");
						}
					}
				}
			}
		}

		for (int i = 0; i < numOfFolds; i++) {
			saveToFile(trainingSets[i], location + "/train" + (i + 1) + ".dat");
			saveToFile(testingSets[i], location + "/test" + (i + 1) + ".dat");
			System.out.println("Training Set for fold " + (i + 1));
			System.out.print(trainingSets[i]);
			System.out.println("Testing Set for fold " + (i + 1));
			System.out.print(testingSets[i]);
		}

	}

	public void saveToFile(String text, String filename) {
		FileOutputStream fstream;
		File file = new File(filename);
		try {
			fstream = new FileOutputStream(file);
			fstream.write(text.getBytes());
			fstream.close();
		} catch (FileNotFoundException e) {
			System.out.println("Attempt to write to " + file.toString()
					+ " failed.");
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}
	
	public static void main(String args[]){
		CrossValidator cv = new CrossValidator();
		cv.loadFiles("C:\\Users\\Amos\\workspace\\Data\\selected\\933\\933b.csv", false);
		cv.crossValidate(5, "C:\\Users\\Amos\\workspace\\Data\\selected\\933\\cv");
	}
}
